package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"gopkg.in/yaml.v2"
	"io/fs"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

func main() {
	root := "./api"
	apiMapSli := []map[string][]string{}
	err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			fmt.Printf("Error accessing path %q: %v\n", path, err)
			return err
		}
		// 如果是目录，打印目录路径
		if d.IsDir() {
			//fmt.Printf("Directory: %s\n", path)
		} else {
			// 如果是文件，打印文件路径
			//fmt.Printf("File: %s\n", path)
			// 如果是 .proto 文件，读取并处理
			if strings.HasSuffix(path, ".proto") {
				apiMapSli = append(apiMapSli, fileReader(path))
			}
		}
		if !d.IsDir() && strings.HasSuffix(path, ".proto") {
			apiMapSli = append(apiMapSli, fileReader(path))
		}
		return nil
	})

	if err != nil {
		fmt.Printf("WalkDir error: %v\n", err)
	}

	//for k, v := range apiMapSli {
	//	fmt.Printf("key: %s\t, val: %s\n", k, v)
	//}
	for _, v := range apiMapSli {
		//fillJson(v, "./openapi.json")
		fillYaml(v, "openapi.yaml")
	}
}

func fileReader(filepath string) map[string][]string {
	m := make(map[string][]string)
	file, err := os.Open(filepath)
	if err != nil {
		log.Fatalf("Error opening file: %v", err)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	tagPattern := regexp.MustCompile(`//\s*@Tags\s+([a-zA-Z0-9_]+(?:\$\$?[a-zA-Z0-9_]+)*)`)
	apiPattern := regexp.MustCompile(`(?i)(get|post|put|delete|patch|options|head):\s*"([^"]+)"`)

	var curTags []string
	for scanner.Scan() {
		line := scanner.Text()
		matchTag := tagPattern.FindStringSubmatch(line)
		matchApi := apiPattern.FindStringSubmatch(line)

		if matchTag != nil {
			tags := strings.Split(matchTag[1], "$")
			curTags = append(curTags, tags...)
		}

		if matchApi != nil && len(curTags) > 0 {
			m[matchApi[0]] = append([]string{}, curTags...)
			curTags = curTags[:0]
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Println("Error reading file:", err)
	}
	return m
}

type OpenAPIS struct {
	OpenAPI    string                                       `json:"openapi"`
	Info       struct{ Title, Version string }              `json:"info"`
	Paths      map[string]map[string]map[string]interface{} `json:"paths"`
	Components struct {
		Schemas map[string]interface{} `json:"schemas"`
	} `json:"components"`
}

func fillJson(midMap map[string][]string, openapiFP string) {
	file, err := os.Open(openapiFP)
	if err != nil {
		fmt.Printf("Failed to open OpenAPI file: %v\n", err)
		return
	}
	defer file.Close()

	data, err := ioutil.ReadAll(file)
	if err != nil {
		fmt.Printf("Failed to read OpenAPI file: %v\n", err)
		return
	}

	// 处理 UTF-8 BOM
	data = removeBOM(data)

	var openapis OpenAPIS
	if err = json.Unmarshal(data, &openapis); err != nil {
		fmt.Printf("Failed to unmarshal OpenAPI JSON: %v\n", err)
		return
	}

	for key, vals := range midMap {
		keySli := strings.Split(key, ":")
		apiPath := strings.Trim(strings.TrimSpace(keySli[1]), `"`)
		method := strings.TrimSpace(strings.ToLower(keySli[0]))

		if _, ok := openapis.Paths[apiPath][method]; ok {
			if tags, ok := openapis.Paths[apiPath][method]["tags"].([]interface{}); ok {
				existingTags := map[string]bool{}
				for _, tag := range tags {
					if t, ok := tag.(string); ok {
						existingTags[t] = true
					}
				}
				for _, val := range vals {
					if !existingTags[val] {
						tags = append(tags, val)
					}
				}
				openapis.Paths[apiPath][method]["tags"] = tags
			} else {
				openapis.Paths[apiPath][method]["tags"] = vals
			}
		}
	}

	output, err := json.MarshalIndent(openapis, "", "  ")
	if err != nil {
		fmt.Printf("Failed to marshal updated OpenAPI JSON: %v\n", err)
		return
	}

	err = ioutil.WriteFile(openapiFP, output, 0644)
	if err != nil {
		fmt.Printf("Failed to write updated OpenAPI file: %v\n", err)
		return
	}
}

type OpenAPIYml struct {
	OpenAPI    string                                       `yaml:"openapi"`
	Info       struct{ Title, Version string }              `yaml:"info"`
	Paths      map[string]map[string]map[string]interface{} `yaml:"paths"`
	Components struct {
		Schemas map[string]interface{} `yaml:"schemas"`
	} `yaml:"components"`
}

func fillYaml(midMap map[string][]string, openapiFP string) {
	file, err := ioutil.ReadFile(openapiFP)
	if err != nil {
		fmt.Printf("Failed to open OpenAPI file: %v\n", err)
		return
	}

	var openapis OpenAPIYml
	if err = yaml.Unmarshal(file, &openapis); err != nil {
		fmt.Printf("Failed to unmarshal OpenAPI YAML: %v\n", err)
		return
	}

	for key, vals := range midMap {
		keySli := strings.Split(key, ":")
		apiPath := strings.Trim(strings.TrimSpace(keySli[1]), `"`)
		method := strings.TrimSpace(strings.ToLower(keySli[0]))

		if _, ok := openapis.Paths[apiPath][method]; ok {
			if tags, ok := openapis.Paths[apiPath][method]["tags"].([]interface{}); ok {
				existingTags := map[string]bool{}
				for _, tag := range tags {
					if t, ok := tag.(string); ok {
						existingTags[t] = true
					}
				}
				for _, val := range vals {
					if !existingTags[val] {
						tags = append(tags, val)
					}
				}
				openapis.Paths[apiPath][method]["tags"] = tags
			} else {
				openapis.Paths[apiPath][method]["tags"] = vals
			}
		}
	}

	output, err := yaml.Marshal(&openapis)
	if err != nil {
		fmt.Printf("Failed to marshal updated OpenAPI YAML: %v\n", err)
		return
	}

	err = ioutil.WriteFile(openapiFP, output, 0644)
	if err != nil {
		fmt.Printf("Failed to write updated OpenAPI file: %v\n", err)
		return
	}
}

func removeBOM(data []byte) []byte {
	if len(data) >= 3 && data[0] == 0xEF && data[1] == 0xBB && data[2] == 0xBF {
		return data[3:]
	}
	return data
}
